'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 5_study_create_xml_agent.py
* @Time: 2025/11/5
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain.agents import AgentExecutor, create_xml_agent
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

dotenv.load_dotenv()

class GoogleSerperSchema(BaseModel):
    query:str = Field(description="执行谷歌搜索的查询语句")

google_serper = GoogleSerperRun(
    name = "google_serper",
    description = (
       "一个低成本的谷歌搜索工具"
       "当你想要回答有关问题的时候,可以调用该工具"
       "该工具的输入是搜索查询语句"),
        api_wrapper=GoogleSerperAPIWrapper()
)


dalle = OpenAIDALLEImageGenerationTool(
    name = "openai-dalle",
    api_wrapper=DallEAPIWrapper(model="dall-e-3")
)

tools = [google_serper, dalle]

prompt = ChatPromptTemplate.from_messages([
    ("human", """You are a helpful assistant. Help the user answer any questions.

You have access to the following tools:

{tools}

In order to use a tool, you can use <tool></tool> and <tool_input></tool_input> tags. You will then get back a response in the form <observation></observation>
For example, if you have a tool called 'search' that could run a google search, in order to search for the weather in SF you would respond:

<tool>search</tool><tool_input>weather in SF</tool_input>
<observation>64 degrees</observation>

When you are done, respond with a final answer between <final_answer></final_answer>. For example:

<final_answer>The weather in SF is 64 degrees</final_answer>

Begin!

Previous Conversation:
{chat_history}

Question: {input}
{agent_scratchpad}"""),
])

# 创建大语言模型
llm = ChatOpenAI(model="gpt-4o-mini")

# 创建agent与agent执行者
agent = create_xml_agent(
    prompt=prompt,
    llm=llm,
    tools=tools,
)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

print(agent_executor.invoke({"input": "你好最近台湾有什么新闻？", "chat_history": ""}))
